import numpy as np


class TSP:
    def __init__(self, distance_matrix, start_city, end_city, num_elite_ants):
        self.distance_matrix = distance_matrix
        self.start_city = start_city
        self.end_city = end_city
        self.unvisited_cities = [city for city in range(len(distance_matrix)) if
                                 city != start_city and city != end_city]
        self.num_elite_ants = num_elite_ants


def aco_tsp(distance_matrix, start_city, end_city, num_ants, num_iterations, alpha, beta, rho, Q, num_elite_ants):
    tsp = TSP(distance_matrix, start_city, end_city, num_elite_ants)
    num_cities = len(distance_matrix)
    # 初始化所有路径信息素大小为1
    pheromone_matrix = np.ones((num_cities, num_cities))

    # 初始化数据
    best_tour = None
    best_tour_length = float('inf')

    display_data = dict()
    display_data["max_pheromone"] = 1

    # 清完善代码，获取正确的best_tour,best_tour_length,display_data-------------------------------------------------------------------------

    best_tour = [1, 0, 2]
    best_tour_length = 400.0
    display_data = {
        'max_pheromone': 19.0,  # 最大信息素浓度
        'iteration_results':
            [
                {
                    'iteration': 1,
                    'pheromone': {'edge-1143': 4.0, 'edge-4': 4.0, 'edge-7': 4.0},  # 各边信息素浓度
                    'ants': [
                        {
                            'path': {
                                'node': [3, 5, 0, 5, 1],  # 经过的节点完整序号
                                'edge': [['edge-4', 'edge-1143'], ['edge-1143', 'edge-7']],  # 该路径经过的边
                                'path': [3, 0, 1]  # 节点5不属于起点终点，也不属于待巡检的点，算法中可以不考虑该点，直接考虑起点终点和待巡检点之间的关系即可，若这种中两点不可直达，用最短路径连接
                            },  # 该蚂蚁本次迭代下的路径信息
                            'cost': 400.0,  # 该蚂蚁本次迭代下的路径长度代价
                            'probabilities': [
                                [
                                    {
                                        'p': 1.0,  # 概率大小
                                        'edge': (3, 0),  # 边，表示从节点3 -> 节点0这条边
                                        'total_pheromone': 1.0,  # 该边当前的信息素
                                        'selected_p': 0.8395780678698624,  # 使用轮盘赌算法时，随机产生的数值
                                        'len': 200.0,  # 该边的长度代价
                                        'edges_id': [3, 5, 0],  # 边两端的完整路径（拓扑图中，从节点3到节点0，最短还需要经过节点5才能到达）
                                        'edges': ['edge-4', 'edge-1143']  # 节点3 -> 节点0需要经过的边
                                    }
                                ]
                            ],  # 该蚂蚁在本次迭代过程中，选择各边时的概率信息，注意当蚂蚁仅剩终点未到达时，无需计算选择到终点的边的概率
                            'pheromone_change': [' 5.00 = 2 * 1000 /  400.00 '],  # 迭代过程中所有边的信息素变化计算公式
                            'is_elite': True,  # 该蚂蚁是否为精英蚂蚁
                            'r': 0.5,  # 信息素挥发因子
                            'pheromone': 5.0  # 本次迭代所产生的总信息素
                        },
                        {
                            'path': {'node': [3, 5, 0, 5, 1], 'edge': [['edge-4', 'edge-1143'], ['edge-1143', 'edge-7']], 'path': [3, 0, 1]}, 'cost': 400.0, 'probabilities': [[{'p': 1.0, 'edge': (3, 0), 'total_pheromone': 1.0,  'selected_p': 0.01831004941396941, 'len': 200.0,'edges_id': [3, 5, 0],  'edges': ['edge-4', 'edge-1143']}]], 'pheromone_change': [' 2.50 = 1000 /  400.00 '], 'is_elite': False, 'r': 0.5,  'pheromone': 2.5
                        }  # 第二只蚂蚁在本次迭代过程中的信息
                    ],
                    'best': {'path': [1, 0, 2], 'cost': 400.0}  # 本次迭代所有蚂蚁中花费最少的路径信息
                },
                {
                    'iteration': 2, 'pheromone': {'edge-1143': 9.5, 'edge-4': 5.75, 'edge-7': 5.75}, 'ants': [{'path': {'node': [3, 5, 0, 5, 1], 'edge': [['edge-4', 'edge-1143'], ['edge-1143', 'edge-7']], 'path': [3, 0, 1]}, 'cost': 400.0, 'probabilities': [[{'p': 1.0, 'edge': (3, 0), 'total_pheromone': 4.25, 'selected_p': 0.5602920382958814, 'len': 200.0, 'edges_id': [3, 5, 0], 'edges': ['edge-4', 'edge-1143']}]], 'pheromone_change': [' 5.00 = 2 * 1000 /  400.00 '], 'is_elite': True, 'r': 0.5,  'pheromone': 5.0}, {'path': {'node': [3, 5, 0, 5, 1], 'edge': [['edge-4', 'edge-1143'], ['edge-1143', 'edge-7']], 'path': [3, 0, 1]}, 'cost': 400.0, 'probabilities': [[{'p': 1.0, 'edge': (3, 0),  'total_pheromone': 4.25, 'selected_p': 0.6231788470222571, 'len': 200.0, 'edges_id': [3, 5, 0], 'edges': ['edge-4', 'edge-1143']}]], 'pheromone_change': [' 2.50 = 1000 /  400.00 '], 'is_elite': False, 'r': 0.5, 'pheromone': 2.5}], 'best': {'path': [1, 0, 2], 'cost': 400.0}
                }  # 第二次迭代
            ],
        'pathMap': {(0, 3): (200.0, [0, 5, 3]), (0, 1): (200.0, [0, 5, 1]), (3, 0): (200.0, [3, 5, 0]), (3, 1): (200.0, [3, 5, 1]), (1, 0): (200.0, [1, 5, 0]), (1, 3): (200.0, [1, 5, 3])}}
        # pathMap映射了当只考虑起点、终点和待巡检点时，两点之间最短路径的代价和路径信息。如(0, 3): (200.0, [0, 5, 3])表示从节点0到节点3，最短距离为200，路径为0 -> 5 -> 3

    return best_tour, best_tour_length, display_data
